﻿using System.Security.Cryptography;

namespace Devonline.Security;

/// <summary>
/// Aes 算法提供加解密的数据安全服务
/// </summary>
public class AesSecurityService(SecuritySetting securitySetting) : SecurityService(securitySetting), ISymmetricSecurityService, ISecurityService
{
    protected readonly object _locker = new();
    protected readonly Dictionary<long, KeyIV> _keyIVs = [];
    protected KeyIV? _lastKeyIV = null;

    /// <summary>
    /// 根据 KeyIV 构造 Aes
    /// </summary>
    /// <param name="key">Aes Key</param>
    /// <param name="iv">Aes IV</param>
    /// <returns>Aes 对象</returns>
    protected Aes GetAes(byte[] key, byte[] iv)
    {
        var aes = Aes.Create();
        aes.KeySize = _securitySetting.KeySize;
        aes.Mode = _securitySetting.CipherMode;
        aes.Key = key;
        aes.IV = iv;
        return aes;
    }
    /// <summary>
    /// 使用加解密流转换数据
    /// </summary>
    /// <param name="transform">加解密转换</param>
    /// <param name="data">数据</param>
    /// <returns></returns>
    protected async Task<byte[]> TransformDataAsync(ICryptoTransform transform, byte[] data)
    {
        using var memoryStream = new MemoryStream();
        using var cryptoStream = new CryptoStream(memoryStream, transform, CryptoStreamMode.Write);
        await cryptoStream.WriteAsync(data);
        cryptoStream.Close();
        var transformed = memoryStream.ToArray();
        memoryStream.Close();
        return transformed;
    }

    /// <summary>
    /// 构造 Aes 的 KeyIV
    /// </summary>
    /// <returns>KeyIV 对象</returns>
    public virtual KeyIV GetKeyIV()
    {
        lock (_locker)
        {
            if (_lastKeyIV is null || _lastKeyIV.UsedCount > _securitySetting.UsageCount || _lastKeyIV.CreatedTime.AddMinutes(_securitySetting.ExpireTime) >= DateTime.UtcNow)
            {
                using var aes = Aes.Create();
                _lastKeyIV = new KeyIV(KeyGenerator.GetKey(), aes.Key, aes.IV);
                _keyIVs.SetValue(_lastKeyIV.Id, _lastKeyIV);
            }
        }

        return _lastKeyIV;
    }
    /// <summary>
    /// 根据密文数据携带的 Aes 密钥编号获取 Aes 密钥
    /// </summary>
    /// <param name="keyId">密钥编号</param>
    /// <returns>返回匹配的 Aes 密钥</returns>
    /// <exception cref="KeyNotFoundException"></exception>
    public virtual KeyIV GetKeyIV(long keyId)
    {
        return _keyIVs.GetValue(keyId) ?? throw new KeyNotFoundException("Aes security service Aes Key Id not found!");
    }
    /// <summary>
    /// 添加一个新的 KeyIV
    /// </summary>
    /// <param name="keyId">密钥编号</param>
    /// <param name="keyIV">KeyIV 对象</param>
    public virtual void SetKeyIV(long keyId, KeyIV keyIV) => _keyIVs.SetValue(keyIV.Id, keyIV);

    /// <summary>
    /// 使用 Key IV 构造的 Aes 算法加密数据 data
    /// </summary>
    /// <param name="key">Aes Key</param>
    /// <param name="iv">Aes IV</param>
    /// <param name="data">待加密的数据</param>
    /// <returns>加密后的数据</returns>
    public virtual async Task<byte[]> EncryptAsync(byte[] key, byte[] iv, byte[] data)
    {
        //1. 根据 Key IV 构造 Aes
        using var aes = GetAes(key, iv);

        //2. 根据 Aes 构造加密转换器
        using var transform = aes.CreateEncryptor();

        //3. 加密数据
        return await TransformDataAsync(transform, data);
    }
    /// <summary>
    /// 使用 Key IV 构造的 Aes 算法解密密文数据 data
    /// </summary>
    /// <param name="key">Aes Key</param>
    /// <param name="iv">Aes IV</param>
    /// <param name="data">待解密的数据</param>
    /// <returns>解密后的数据</returns>
    public virtual async Task<byte[]> DecryptAsync(byte[] key, byte[] iv, byte[] data)
    {
        //1. 根据 Key IV 构造 Aes
        using var aes = GetAes(key, iv);

        //2. 根据 Aes 构造解密转换器
        using var transform = aes.CreateDecryptor(aes.Key, aes.IV);

        //3. 解密数据
        return await TransformDataAsync(transform, data);
    }

    /// <summary>
    /// 使用 Key IV 构造的 Aes 算法加密数据 plainText
    /// </summary>
    /// <param name="key">Aes Key</param>
    /// <param name="iv">Aes IV</param>
    /// <param name="plainText">待加密的字符串</param>
    /// <returns>Base64 编码的密文字符串</returns>
    public virtual async Task<string> EncryptAsync(byte[] key, byte[] iv, string plainText)
    {
        using var aes = GetAes(key, iv);
        using var transform = aes.CreateEncryptor(aes.Key, aes.IV);
        using var memoryStream = new MemoryStream();
        using var cryptoStream = new CryptoStream(memoryStream, transform, CryptoStreamMode.Write);
        using var streamWriter = new StreamWriter(cryptoStream);
        await streamWriter.WriteAsync(plainText);
        var data = memoryStream.ToArray();
        streamWriter.Close();
        cryptoStream.Close();
        memoryStream.Close();
        return Convert.ToBase64String(data);
    }
    /// <summary>
    /// 使用 Key IV 构造的 Aes 算法解密密文数据 cipherText
    /// </summary>
    /// <param name="key">Aes Key</param>
    /// <param name="iv">Aes IV</param>
    /// <param name="cipherText">Base64 编码的密文字符串</param>
    /// <returns>解密后的数据</returns>
    public virtual async Task<string> DecryptAsync(byte[] key, byte[] iv, string cipherText)
    {
        var data = Convert.FromBase64String(cipherText);
        using var aes = GetAes(key, iv);
        using var transform = aes.CreateDecryptor(aes.Key, aes.IV);
        using var memoryStream = new MemoryStream(data);
        using var cryptoStream = new CryptoStream(memoryStream, transform, CryptoStreamMode.Read);
        using var streamReader = new StreamReader(cryptoStream);
        var text = await streamReader.ReadToEndAsync();
        streamReader.Close();
        cryptoStream.Close();
        memoryStream.Close();
        return text;
    }

    /// <summary>
    /// 使用 Aes 算法加密明文数据
    /// </summary>
    /// <param name="data">明文数据</param>
    /// <returns>带密钥编号的数据密文</returns>
    public override async Task<byte[]> EncryptAsync(byte[] data)
    {
        var keyIV = GetKeyIV();
        var encrypted = await EncryptAsync(keyIV.Key, keyIV.IV, data);
        keyIV.Used();
        return new KeyData(keyIV.Id, encrypted).Value;
    }
    /// <summary>
    /// 使用 Aes 算法解密密文数据
    /// </summary>
    /// <param name="data">带密钥编号的数据密文</param>
    /// <returns>解密后的明文数据</returns>
    public override async Task<byte[]> DecryptAsync(byte[] data)
    {
        var keyData = new KeyData(data);
        if (keyData.EncryptionMode != DataEncryptionMode.KeyId)
        {
            throw new ArgumentException("Aes security service decrypt data must have Aes Key Id!");
        }

        var keyIV = GetKeyIV(keyData.KeyId);
        return await DecryptAsync(keyIV.Key, keyIV.IV, keyData.Data);
    }
}